from datetime import datetime, timedelta

from flask_restful import Resource
from flask import current_app, request, g
import random

from flask_restful.reqparse import RequestParser

from celery_tasks.sms.tasks import send_verification_code
from models import db
from models.user import User, UserProfile
from utils import parser
from utils.decorators import set_db_to_write, set_db_to_read, login_required
from . import constants
from utils.limiter import limiter as lmt
from flask_limiter.util import get_remote_address
from utils.jwt_util import generate_jwt, verify_jwt


class SMSVerificationCodeResource(Resource):
    """
    获取短信验证码
    """
    error_message = 'Too many requests.'

    decorators = [
        lmt.limit(constants.LIMIT_SMS_VERIFICATION_CODE_BY_MOBILE,
                  key_func=lambda: request.view_args['mobile'],
                  error_message=error_message),
        lmt.limit(constants.LIMIT_SMS_VERIFICATION_CODE_BY_IP,
                  key_func=get_remote_address,
                  error_message=error_message)
    ]

    def get(self, mobile):
        # 生成验证码
        code = '{:0>6d}'.format(random.randint(0, 999999))
        # 保存到redis中
        current_app.redis_master.setex('app:code:{}'.format(mobile), constants.SMS_VERIFICATION_CODE_EXPIRES, code)
        # 使用celery异步任务 发送短信
        send_verification_code.delay(mobile, code)
        return {'mobile': mobile}


class AuthorizationResource(Resource):
    """
    登录认证
    """
    method_decorators = {
        'post': [set_db_to_write],
        'get': [login_required]
    }

    def _generate_token(self, user_id, with_refresh_token=True):
        """
        生成JWT

        :param user_id: 用户id
        :return: 访问和刷新token
        """
        """生成访问token"""

        # 构建数据
        payload = {'payload': user_id, 'is_refresh': False}
        expiry = datetime.utcnow() + timedelta(current_app.config['JWT_EXPIRY_HOURS'])
        access_token = generate_jwt(payload, expiry)

        """生成刷新token"""
        if with_refresh_token == True:
            # 构建数据
            payload = {'payload': user_id, 'is_refresh': True}
            expiry = datetime.utcnow() + timedelta(current_app.config['JWT_REFRESH_DAYS'])
            refresh_token = generate_jwt(payload, expiry)
        else:
            refresh_token = None

        return access_token, refresh_token

    def post(self):
        """
        登录创建token
        """
        json_parser = RequestParser()
        json_parser.add_argument('mobile', type=parser.mobile, required=True, location='json')
        json_parser.add_argument('code', type=parser.regex(r'^\d{6}$'), required=True, location='json')
        args = json_parser.parse_args()
        mobile = args.mobile
        code = args.code

        # 从redis中获取验证码
        key = 'app:code:{}'.format(mobile)
        try:
            real_code = current_app.redis_master.get(key)
        except ConnectionError as e:  # 如果主数据库连接失败, 再到从数据库中读取
            current_app.logger.error(e)
            real_code = current_app.redis_slave.get(key)

        # try:
        #     current_app.redis_master.delete(key)  # 删除验证码. 每个验证码只能使用一次
        # except ConnectionError as e:
        #     current_app.logger.error(e)

        if not real_code or real_code.decode() != code:
            return {'message': 'Invalid code.'}, 400

        # 查询或保存用户
        user = User.query.filter_by(mobile=mobile).first()

        if user is None:
            # 用户不存在，注册用户
            user_id = current_app.id_worker.get_id()  # 生成分布式id
            user = User(id=user_id, mobile=mobile, name=mobile, last_login=datetime.now())
            db.session.add(user)
            profile = UserProfile(id=user.id)
            db.session.add(profile)
            db.session.commit()

        # TODO 免密码登录
        access_token, refresh_token = self._generate_token(user.id)
        # 201 表示生成成功
        return {'access_token': access_token, 'refresh_token': refresh_token}, 201

    def get(self):  # 测试访问控制
        return {"user_id": g.user_id, "is_refresh": g.is_refresh}

    def put(self):
        """刷新 访问token"""

        # 判断 是否认证通过, 并且发送的是 刷新token
        if g.user_id and g.is_refresh:
            # 生成新的访问token
            access_token, refresh_token = self._generate_token(g.user_id, with_refresh_token=False)
            # 返回数据
            return {'access_token': access_token}
        else:
            return {'message': 'Invalid token'}, 401